#!/usr/bin/env python
# -*- coding=utf-8 -*-
"""
@author: xingwg
@license: (C) Copyright 2020-2025.
@contact: xingweiguo@chinasvt.com
@project: boya-reid
@file: data_augmentation.py
@time: 2020/9/13 0:48
@desc:
"""
import math
import random
import cv2
import numpy as np
from PIL import Image


class RandomErasing(object):
    """Randomly selects a rectangle region in an image and erases its pixels.
        'Random Erasing Data Augmentation' by Zhong et al.
        See https://arxiv.org/pdf/1708.04896.pdf
    Args:
         probability: The probability that the Random Erasing operation will be performed.
         sl: Minimum proportion of erased area against input image.
         sh: Maximum proportion of erased area against input image.
         r1: Minimum aspect ratio of erased area.
         mean: Erasing value.
    """
    def __init__(self, probability=0.5, sl=0.02, sh=0.4, r1=0.3, mean=(0.4914, 0.4822, 0.4465)):
        self.probability = probability
        self.mean = mean
        self.sl = sl
        self.sh = sh
        self.r1 = r1

    def __call__(self, img):
        """

        Args:
            img: torch tensor

        Returns:

        """
        if random.uniform(0, 1) >= self.probability:
            return img

        for attempt in range(100):
            area = img.size()[1] * img.size()[2]

            target_area = random.uniform(self.sl, self.sh) * area
            aspect_ratio = random.uniform(self.r1, 1 / self.r1)

            h = int(round(math.sqrt(target_area * aspect_ratio)))
            w = int(round(math.sqrt(target_area / aspect_ratio)))

            if w < img.size()[2] and h < img.size()[1]:
                x1 = random.randint(0, img.size()[1] - h)
                y1 = random.randint(0, img.size()[2] - w)
                if img.size()[0] == 3:
                    img[0, x1:x1 + h, y1:y1 + w] = self.mean[0]
                    img[1, x1:x1 + h, y1:y1 + w] = self.mean[1]
                    img[2, x1:x1 + h, y1:y1 + w] = self.mean[2]
                else:
                    img[0, x1:x1 + h, y1:y1 + w] = self.mean[0]
                return img

        return img


class RandomShift(object):
    """
    2019NAIC行人重识别大赛冠军数据增强方法
    """
    def __init__(self, prob=0.5, rc_min=0.7, rh_min=0.5, rw_min=0.5, mean=None):
        """
        构造函数
        Args:
            prob: 触发随机平移的概率
            rc_min:
            rh_min:
            rw_min:
        """
        super(RandomShift, self).__init__()
        self.prob = prob
        self.rc_min = rc_min
        self.rh_min = rh_min
        self.rw_min = rw_min
        self.mean = (0, 0, 0) if mean is None else tuple(mean)

    def __call__(self, img):
        """

        Args:
            img: opencv image

        Returns:

        """
        if random.uniform(0, 1) >= self.prob:
            img = img*255.0
            img = img.astype("uint8")
            return Image.fromarray(img)

        rc = random.uniform(self.rc_min, 1.0)

        h, w, c = img.shape

        crop_h = int(h * rc)

        x1 = 0
        y1 = random.randint(0, h - crop_h)
        x2 = w
        y2 = crop_h + y1

        # crop_img = img.crop((x1, y1, x2, y2))
        crop_img = img[y1:y2, x1:x2, :]

        rw = random.uniform(self.rw_min, 1.0)
        rh = random.uniform(self.rh_min, 1. / rc)

        nw, nh = int(rw * w), int(rh * crop_h)

        # resize_img = crop_img.resize((nw, nh))
        resize_img = cv2.resize(crop_img, (nw, nh))

        # bg_img = Image.new("RGB", (w, h), (0, 0, 0))
        bg_img = np.ones((h, w, c), dtype=np.int8) * self.mean

        offset_x = random.randint(0, w - nw)
        offset_y = random.randint(0, h - nh)

        # img = bg_img.paste(resize_img, (shift_x, shift_y, shift_x + nw, shift_y + nh))
        bg_img[offset_y:offset_y + nh, offset_x:offset_x + nw, :] = resize_img

        img = bg_img.copy()
        img = img * 255.0
        img = img.astype("uint8")

        return Image.fromarray(img)


